<?php

namespace App\Console\Commands\model;

use App\Dep\ModelDep;
use App\Dep\SystemDep;
use App\Dep\TrainLogDep;
use App\Enum\CommonEnum;
use App\Enum\TrainEnum;
use GuzzleHttp\Client;
use Illuminate\Console\Command;

class DoTrainCmd extends Command
{
    //打标
    public $signature = 'model:do-train';

    public function handle()
    {
        $dep = new ModelDep();
        $logDep = new TrainLogDep();
        $systemDep = new SystemDep();

        $resSystem = $systemDep->first();
        $client = new Client();

        $response = $client->request('get', $resSystem->tag_domain . '/api/tasks');
        $taskRes = json_decode($response->getBody(), true);
        if(!isset($taskRes['status']) || $taskRes['status']!='success'){
            $this->log('server error',$taskRes);
            return false;
        }
        $this->log('server tasks',$taskRes);

        $taskData = $taskRes['data']['tasks'];
        foreach ($taskData as $task){
            if($task['status'] == 'RUNNING'){
                $this->log('server running',$task);
                return false;
            }
        }

        $resWaitTrain = $dep->firstWaitTrain();
        if (!$resWaitTrain) {
            $this->log('have no train model');
            return false;
        }

        $trainDir = "{$resSystem->train_dir}\\{$resWaitTrain->name}";
        if(file_exists($trainDir)){
            del_dir($trainDir);
        }
        mkdir($trainDir);
        $this->log('recreate train dir',['train_dir'=>$trainDir]);

        $sourceArr = json_decode($resWaitTrain->source,true);

        foreach ($sourceArr as $source){
            $originPath = $resSystem->path . '\\' . '训练模型' . '\\' . $resWaitTrain->uuid . '\\素材图' . '\\' . $source['source'];
            $distPath = "{$trainDir}\\{$source['source']}";
            $txtPath = "{$trainDir}\\{$source['title']}.txt";
            copy($originPath,$distPath);
            file_put_contents($txtPath,$source['origin_tag']);
            $this->log('copy source',['textPath'=>$txtPath]);
        }

        $jsonData = [
            "pretrained_model_name_or_path" => $resWaitTrain->mold,
            "train_data_dir" => $trainDir,
            "resolution" => $resWaitTrain->size,
            "enable_bucket" => true,
            "min_bucket_reso" => 256,
            "max_bucket_reso" => 1024,
            "output_name" => $resWaitTrain->name,
            "output_dir" => $trainDir,
            "save_model_as" => "safetensors",
            "save_every_n_epochs" => 2,
            "max_train_epochs" => 10,
            "train_batch_size" => 1,
            "network_train_unet_only" => false,
            "network_train_text_encoder_only" => false,
            "learning_rate" => 0.0001,
            "unet_lr" => 0.0001,
            "text_encoder_lr" => 0.00001,
            "lr_scheduler" => "cosine_with_restarts",
            "optimizer_type" => "AdamW8bit",
            "lr_scheduler_num_cycles" => 1,
            "network_module" => "networks.lora",
            "network_dim" => 32,
            "network_alpha" => 32,
            "logging_dir" => "./logs",
            "caption_extension" => ".txt",
            "shuffle_caption" => true,
            "keep_tokens" => 0,
            "max_token_length" => 255,
            "seed" => 1337,
            "prior_loss_weight" => 1,
            "clip_skip" => 2,
            "mixed_precision" => "fp16",
            "save_precision" => "fp16",
            "xformers" => true,
            "cache_latents" => true,
            "persistent_data_loader_workers" => true,
            "lr_warmup_steps" => 0
        ];
        $this->log('request server', $jsonData);

        $response = $client->request('post', $resSystem->tag_domain . '/api/run', [
            'json' => $jsonData
        ]);

        $body = $response->getBody();
        $res = json_decode($body, true);
        $this->log('res server', $res);

        if(!isset($res['status']) || $res['status'] != 'success'){
            $this->log('server error',$res);
            $data = [
              'status' => TrainEnum::TRAIN_STATUS_TRAIN_ERR,
            ];
            $dep->edit($resWaitTrain->id, $data);

            $logData = [
                'content' => "训练出错:{$res['message']}",
                'model_id' => $resWaitTrain->id,
            ];
            $logDep->add($logData);
            return false;
        }

        $data = [
            'is_train' => CommonEnum::YES,
            'train_at' => date('Y-m-d H:i:s'),
            'train_num' => $res['message']
        ];
        $dep->edit($resWaitTrain->id, $data);

        $logData = [
            'content' => '开始训练',
            'model_id' => $resWaitTrain->id,
        ];
        $logDep->add($logData);
    }

    private function checkOk($sourceArr)
    {
        foreach ($sourceArr as $source) {
            if (!isset($source['origin_tag'])) {
                return false;
            }
        }
        return true;
    }

    private function log($msg, $context = [])
    {
        $this->info("do-train log: {$msg} " . json_encode($context,JSON_UNESCAPED_UNICODE));
        log_daily("do-train")->info($msg, $context);
    }


}
